#include "stdafx.h"
#include "Registration.h"
#include "MiniBasicImpl.h"
#include <shlwapi.h>

std::wstring CppStringFromIID(REFIID iid)
{
    // string form of guid is fixed at 
    // {c200e360-38c5-11ce-ae62-08002b2b79ef}
    // which requires 39 chars (including NULL)
    // just rounded up to 40 to make it a round figure
    WCHAR stringGuid[40];
    int charsCopied = StringFromGUID2(iid, stringGuid, ARRAYSIZE(stringGuid));
    return std::wstring(stringGuid, charsCopied - 1);
}

LONG CreateInprocSubKey(HKEY hParentKey)
{
    HKEY hInprocKey = NULL;
    LONG result = RegCreateKeyEx(hParentKey, L"InprocServer32", 0, NULL, REG_OPTION_NON_VOLATILE, KEY_SET_VALUE, NULL, &hInprocKey, NULL);
    if(result == ERROR_SUCCESS)
    {
        try
        {
            std::wstring modulePath = MiniBasicImpl::GetFullTypeLibPath();
            std::wstring::size_type lastSlash = modulePath.rfind(L'\\');
            modulePath[lastSlash] = 0;
            result = RegSetValueEx(hInprocKey, NULL, 0, REG_SZ, reinterpret_cast<const BYTE*>(modulePath.c_str()), (lastSlash + 1) * sizeof(std::wstring::value_type));
            if(result == ERROR_SUCCESS)
            {
                static const WCHAR threading[] = L"Both";
                result = RegSetValueEx(hInprocKey, L"ThreadingModel", 0, REG_SZ, reinterpret_cast<const BYTE*>(threading), (ARRAYSIZE(threading) + 1) * sizeof(threading[0]));
            }
        }
        catch(const std::runtime_error&)
        {
            result = ERROR_DLL_INIT_FAILED;
        }
        RegCloseKey(hInprocKey);
    }
    return result;
}

STDAPI CreateCLSIDKey(const std::wstring& clsidPath, const std::wstring& title, const std::wstring& baseProgId, const std::wstring& verProgId)
{
    HKEY hClsidKey = NULL;
    LONG result = RegCreateKeyEx(HKEY_CLASSES_ROOT, clsidPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_SET_VALUE, NULL, &hClsidKey, NULL);
    if(result == ERROR_SUCCESS)
    {
        result = RegSetValueEx(hClsidKey, NULL, 0, REG_SZ, reinterpret_cast<const BYTE*>(title.c_str()), (title.size() + 1) * sizeof(std::wstring::value_type));
        if(result == ERROR_SUCCESS)
        {
            result = CreateInprocSubKey(hClsidKey);
            if(result == ERROR_SUCCESS)
            {
                result = RegSetValueEx(hClsidKey, L"ProgID", 0, REG_SZ, reinterpret_cast<const BYTE*>(baseProgId.c_str()), (baseProgId.size() + 1) * sizeof(std::wstring::value_type));
                if(result == ERROR_SUCCESS)
                {
                    result = RegSetValueEx(hClsidKey, L"VersionIndependentProgID", 0, REG_SZ, reinterpret_cast<const BYTE*>(verProgId.c_str()), (verProgId.size() + 1) * sizeof(std::wstring::value_type));
                }
            }
        }
        RegCloseKey(hClsidKey);
    }
    return result != ERROR_SUCCESS ? HRESULT_FROM_WIN32(result) : S_OK;
}

STDAPI CreateRegistryPaths(const std::wstring& baseProgId, REFIID clsid, ITypeLib* pTypeLib, std::wstring& clsidPath, std::wstring& versionedProgId)
{
    (clsidPath = L"CLSID\\") += CppStringFromIID(clsid);
    WORD curVer = 0;

    ITypeInfo* pInfo = NULL;
    HRESULT hr = pTypeLib->GetTypeInfoOfGuid(clsid, &pInfo);
    if(SUCCEEDED(hr))
    {
        TYPEATTR* pTypeInfo = NULL;
        if(SUCCEEDED(hr = pInfo->GetTypeAttr(&pTypeInfo)))
        {
            curVer = pTypeInfo->wMajorVerNum;
            pInfo->ReleaseTypeAttr(pTypeInfo);
        }
        pInfo->Release();
    }
    if(SUCCEEDED(hr))
    {
        std::wostringstream verProgIdCreator;
        verProgIdCreator << baseProgId << L'.' << curVer;
        versionedProgId = verProgIdCreator.str();
    }
    return hr;
}

STDAPI CreateProgIDKeys(const std::wstring& baseProgId, const std::wstring& verProgId, REFIID clsid)
{
    HKEY hBaseKey = NULL;
    LONG status = RegCreateKeyEx(HKEY_CLASSES_ROOT, baseProgId.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_SET_VALUE, NULL, &hBaseKey, NULL);
    const std::wstring& clsidAsString = CppStringFromIID(clsid);
    if(status == ERROR_SUCCESS)
    {
        status = RegSetValueEx(hBaseKey, L"CLSID", 0, REG_SZ, reinterpret_cast<const BYTE*>(clsidAsString.c_str()), (clsidAsString.size() + 1) * sizeof(std::wstring::value_type));
        RegSetValueEx(hBaseKey, L"CurVer", 0, REG_SZ, reinterpret_cast<const BYTE*>(verProgId.c_str()), (verProgId.size() + 1) * sizeof(std::wstring::value_type));
        RegCloseKey(hBaseKey);
        if(status == ERROR_SUCCESS)
        {
            HKEY hVerKey = NULL;
            status = RegCreateKeyEx(HKEY_CLASSES_ROOT, verProgId.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_SET_VALUE, NULL, &hVerKey, NULL);
            if(status == ERROR_SUCCESS)
            {
                status = RegSetValueEx(hVerKey, L"CLSID", 0, REG_SZ, reinterpret_cast<const BYTE*>(clsidAsString.c_str()), (clsidAsString.size() + 1) * sizeof(std::wstring::value_type));
                RegCloseKey(hVerKey);
            }
        }
    }
    return status != ERROR_SUCCESS ? HRESULT_FROM_WIN32(status) : S_OK;
}

STDAPI CreateRegistryEntries(REFIID clsid, const std::wstring& baseProgId, ITypeLib* pTypeLib, const WCHAR* desc)
{
    std::wstring verProgId, clsidRegPath;
    HRESULT hr = CreateRegistryPaths(baseProgId, clsid, pTypeLib, clsidRegPath, verProgId);
    if(SUCCEEDED(hr))
    {
        if(!(
                SUCCEEDED(
                    hr = CreateCLSIDKey(clsidRegPath, desc ? desc : L"", baseProgId, verProgId)
                ) && 
                SUCCEEDED(
                    hr = CreateProgIDKeys(baseProgId, verProgId, clsid)
                )
            )
        )
        {
            RemoveRegistryEntries(clsid, baseProgId, pTypeLib);
        }
    }
    return hr;
}

STDAPI RemoveRegistryEntries(REFIID clsid, const std::wstring& baseProgId, ITypeLib* pTypeLib)
{
    std::wstring verProgId, clsidRegPath;
    HRESULT hr = CreateRegistryPaths(baseProgId, clsid, pTypeLib, clsidRegPath, verProgId);
    if(SUCCEEDED(hr))
    {
        // remove any and all subkeys
        SHDeleteKey(HKEY_CLASSES_ROOT, clsidRegPath.c_str());
        SHDeleteKey(HKEY_CLASSES_ROOT, verProgId.c_str());
        SHDeleteKey(HKEY_CLASSES_ROOT, baseProgId.c_str());
    }
    return hr;
}